#!/usr/bin/env python3
"""
清空向量化数据脚本
删除Milvus集合中的所有已向量化内容
"""

import json
import sys
import os
import logging
from pathlib import Path

# 添加当前目录到Python路径
sys.path.append(str(Path(__file__).parent))

try:
    from pymilvus import connections, Collection, utility, MilvusException
except ImportError:
    print("请安装必要的依赖包：")
    print("pip install pymilvus")
    exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class VectorDataCleaner:
    """向量数据清理器"""
    
    def __init__(self, milvus_host: str = "localhost", milvus_port: int = 19530):
        """
        初始化清理器
        
        Args:
            milvus_host: Milvus主机地址
            milvus_port: Milvus端口
        """
        self.milvus_host = milvus_host
        self.milvus_port = milvus_port
        self._connect_milvus()
    
    def _connect_milvus(self):
        """连接到Milvus数据库"""
        try:
            connections.connect(
                alias="default",
                host=self.milvus_host,
                port=self.milvus_port
            )
            logger.info(f"成功连接到Milvus: {self.milvus_host}:{self.milvus_port}")
        except Exception as e:
            logger.error(f"连接Milvus失败: {e}")
            raise
    
    def list_collections(self) -> list:
        """
        列出所有集合
        
        Returns:
            集合名称列表
        """
        try:
            collections = utility.list_collections()
            logger.info(f"找到 {len(collections)} 个集合: {collections}")
            return collections
        except Exception as e:
            logger.error(f"获取集合列表失败: {e}")
            return []
    
    def get_collection_info(self, collection_name: str) -> dict:
        """
        获取集合信息
        
        Args:
            collection_name: 集合名称
            
        Returns:
            集合信息字典
        """
        try:
            if not utility.has_collection(collection_name):
                return {"exists": False}
            
            collection = Collection(collection_name)
            collection.load()
            
            info = {
                "exists": True,
                "name": collection_name,
                "entities_count": collection.num_entities,
                "description": collection.description if hasattr(collection, 'description') else "无描述"
            }
            
            return info
        except Exception as e:
            logger.error(f"获取集合 {collection_name} 信息失败: {e}")
            return {"exists": False, "error": str(e)}
    
    def clear_collection_data(self, collection_name: str) -> bool:
        """
        清空集合中的所有数据（保留集合结构）
        
        Args:
            collection_name: 集合名称
            
        Returns:
            是否成功
        """
        try:
            if not utility.has_collection(collection_name):
                logger.warning(f"集合 {collection_name} 不存在")
                return False
            
            collection = Collection(collection_name)
            
            # 获取清理前的实体数量
            before_count = collection.num_entities
            logger.info(f"清理前集合 {collection_name} 包含 {before_count} 个实体")
            
            if before_count == 0:
                logger.info(f"集合 {collection_name} 已经是空的")
                return True
            
            # 删除所有数据 - 使用表达式删除所有实体
            # 对于字符串主键，我们需要获取所有ID然后删除
            collection.load()
            
            # 查询所有ID
            results = collection.query(
                expr="",  # 空表达式匹配所有记录
                output_fields=["id"],
                limit=before_count + 1000  # 确保获取所有记录
            )
            
            if results:
                # 批量删除
                batch_size = 1000
                total_deleted = 0
                
                for i in range(0, len(results), batch_size):
                    batch = results[i:i + batch_size]
                    ids_to_delete = [item["id"] for item in batch]
                    
                    # 构建删除表达式
                    id_list = "', '".join(ids_to_delete)
                    delete_expr = f"id in ['{id_list}']"
                    
                    collection.delete(delete_expr)
                    total_deleted += len(ids_to_delete)
                    logger.info(f"已删除 {total_deleted}/{len(results)} 个实体")
                
                # 刷新集合以确保删除生效
                collection.flush()
                
                # 检查删除结果
                after_count = collection.num_entities
                logger.info(f"清理后集合 {collection_name} 包含 {after_count} 个实体")
                
                if after_count == 0:
                    logger.info(f"成功清空集合 {collection_name}")
                    return True
                else:
                    logger.warning(f"集合 {collection_name} 仍包含 {after_count} 个实体")
                    return False
            else:
                logger.info(f"集合 {collection_name} 查询结果为空，可能已经被清空")
                return True
                
        except Exception as e:
            logger.error(f"清空集合 {collection_name} 失败: {e}")
            return False
    
    def drop_collection(self, collection_name: str) -> bool:
        """
        完全删除集合（包括结构）
        
        Args:
            collection_name: 集合名称
            
        Returns:
            是否成功
        """
        try:
            if not utility.has_collection(collection_name):
                logger.warning(f"集合 {collection_name} 不存在")
                return False
            
            # 删除集合
            utility.drop_collection(collection_name)
            logger.info(f"成功删除集合 {collection_name}")
            return True
            
        except Exception as e:
            logger.error(f"删除集合 {collection_name} 失败: {e}")
            return False
    
    def clear_all_collections(self, drop_collections: bool = False) -> dict:
        """
        清空所有集合的数据
        
        Args:
            drop_collections: 是否删除集合结构（默认False，只清空数据）
            
        Returns:
            操作结果统计
        """
        collections = self.list_collections()
        results = {
            "total": len(collections),
            "success": 0,
            "failed": 0,
            "details": []
        }
        
        for collection_name in collections:
            try:
                # 获取集合信息
                info = self.get_collection_info(collection_name)
                
                if not info.get("exists", False):
                    results["details"].append({
                        "collection": collection_name,
                        "status": "不存在",
                        "action": "跳过"
                    })
                    continue
                
                entities_before = info.get("entities_count", 0)
                
                if drop_collections:
                    # 完全删除集合
                    if self.drop_collection(collection_name):
                        results["success"] += 1
                        results["details"].append({
                            "collection": collection_name,
                            "status": "成功删除",
                            "action": "删除集合",
                            "entities_before": entities_before
                        })
                    else:
                        results["failed"] += 1
                        results["details"].append({
                            "collection": collection_name,
                            "status": "删除失败",
                            "action": "删除集合",
                            "entities_before": entities_before
                        })
                else:
                    # 只清空数据
                    if self.clear_collection_data(collection_name):
                        results["success"] += 1
                        results["details"].append({
                            "collection": collection_name,
                            "status": "成功清空",
                            "action": "清空数据",
                            "entities_before": entities_before,
                            "entities_after": 0
                        })
                    else:
                        results["failed"] += 1
                        results["details"].append({
                            "collection": collection_name,
                            "status": "清空失败",
                            "action": "清空数据",
                            "entities_before": entities_before
                        })
                        
            except Exception as e:
                results["failed"] += 1
                results["details"].append({
                    "collection": collection_name,
                    "status": f"操作失败: {e}",
                    "action": "删除集合" if drop_collections else "清空数据"
                })
        
        return results

def load_config(config_path: str = "config.json") -> dict:
    """加载配置文件"""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        logger.warning(f"配置文件 {config_path} 不存在，使用默认配置")
        return {
            "milvus": {
                "host": "localhost",
                "port": 19530,
                "collection_name": "md_documents"
            }
        }
    except json.JSONDecodeError as e:
        logger.error(f"配置文件格式错误: {e}")
        return None

def print_results(results: dict):
    """打印操作结果"""
    print("\n" + "="*60)
    print("清理结果汇总")
    print("="*60)
    print(f"总集合数: {results['total']}")
    print(f"成功处理: {results['success']}")
    print(f"处理失败: {results['failed']}")
    
    if results['details']:
        print("\n详细结果:")
        print("-" * 60)
        for detail in results['details']:
            print(f"集合: {detail['collection']}")
            print(f"  操作: {detail['action']}")
            print(f"  状态: {detail['status']}")
            if 'entities_before' in detail:
                print(f"  清理前实体数: {detail['entities_before']}")
            if 'entities_after' in detail:
                print(f"  清理后实体数: {detail['entities_after']}")
            print()

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="清空向量化数据工具")
    parser.add_argument("-c", "--config", default="config.json", help="配置文件路径")
    parser.add_argument("--collection", help="指定要清理的集合名称（默认清理所有集合）")
    parser.add_argument("--drop", action="store_true", help="完全删除集合而不是只清空数据")
    parser.add_argument("--list", action="store_true", help="只列出集合信息，不执行清理")
    parser.add_argument("--force", action="store_true", help="强制执行，不询问确认")
    
    args = parser.parse_args()
    
    # 加载配置
    config = load_config(args.config)
    if not config:
        return 1
    
    # 初始化清理器
    try:
        cleaner = VectorDataCleaner(
            milvus_host=config["milvus"]["host"],
            milvus_port=config["milvus"]["port"]
        )
    except Exception as e:
        logger.error(f"初始化失败: {e}")
        return 1
    
    # 获取集合列表
    collections = cleaner.list_collections()
    
    if not collections:
        print("没有找到任何集合")
        return 0
    
    # 如果只是列出信息
    if args.list:
        print("\n" + "="*60)
        print("集合信息列表")
        print("="*60)
        for collection_name in collections:
            info = cleaner.get_collection_info(collection_name)
            if info.get("exists", False):
                print(f"集合: {collection_name}")
                print(f"  实体数量: {info.get('entities_count', 0)}")
                print(f"  描述: {info.get('description', '无')}")
                print()
            else:
                print(f"集合: {collection_name} (无法获取信息)")
        return 0
    
    # 确定要处理的集合
    if args.collection:
        if args.collection not in collections:
            logger.error(f"指定的集合 {args.collection} 不存在")
            return 1
        target_collections = [args.collection]
    else:
        target_collections = collections
    
    # 显示将要执行的操作
    action = "删除" if args.drop else "清空数据"
    print(f"\n将要{action}以下集合:")
    for collection_name in target_collections:
        info = cleaner.get_collection_info(collection_name)
        entities_count = info.get('entities_count', 0) if info.get('exists', False) else 0
        print(f"  - {collection_name} (包含 {entities_count} 个实体)")
    
    # 确认操作
    if not args.force:
        confirm = input(f"\n确认要{action}这些集合吗？(y/N): ").strip().lower()
        if confirm not in ['y', 'yes']:
            print("操作已取消")
            return 0
    
    # 执行清理
    print(f"\n开始{action}...")
    
    if args.collection:
        # 处理单个集合
        if args.drop:
            success = cleaner.drop_collection(args.collection)
        else:
            success = cleaner.clear_collection_data(args.collection)
        
        if success:
            print(f"成功{action}集合 {args.collection}")
        else:
            print(f"失败{action}集合 {args.collection}")
            return 1
    else:
        # 处理所有集合
        results = cleaner.clear_all_collections(drop_collections=args.drop)
        print_results(results)
        
        if results['failed'] > 0:
            return 1
    
    print(f"\n{action}操作完成!")
    return 0

if __name__ == "__main__":
    exit(main())
